Schrodinger equation in 1D

In [17]:

    

%matplotlib inline
from pylab import *
from numpy import *


    

In [18]:

    

#physical constants
hbar = 1.0
m = 1.0


    

In [19]:

    

#discretization parameters
nx = 100   # number of unknown grid points in spatial direction
SC = 1.0  # stability criterion SC = 2*D*dt/dx^2, for FTCS should be SC <= 1

def schr_init(maxx, maxt, maxV, nx, SC):
    #choose time step according to CFL condition
    dx = maxx/(nx+1)
    dt = hbar/(2*hbar**2/(m*dx**2)+maxV)*SC
    #dt = SC*dx**2*m/hbar  # for V = 0
    nt = int(maxt/dt)+1
    
    #define array for storing the solution
    Ur = zeros((nt, nx+2))
    Ui = zeros((nt, nx+2))
    
    x = arange(nx+2)*dx
    t = arange(nt)*dt
    return Ur, Ui, dx, dt, x, t


    

In [20]:

    

def schr_solve(Ur, Ui, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        Ui[it+1,xint] = Ui[it,xint] + hbar*dt/(2*m*dx**2)*(Ur[it,xint+1] - 2*Ur[it, xint] + Ur[it, xint-1])\
            - dt/hbar*V[xint]*Ur[it,xint]
        Ur[it+1,xint] = Ur[it,xint] - hbar*dt/(2*m*dx**2)*(Ui[it+1,xint+1] - 2*Ui[it+1, xint] + Ui[it+1, xint-1])\
            + dt/hbar*V[xint]*Ui[it+1,xint]


    

In [21]:

    

maxx = 2.
maxt = 0.05
nx = 200
Ur, Ui, dx, dt, x, t = schr_init(maxx, maxt, 0, nx, 1.0)
V = zeros_like(x)


    

In [22]:

    

sigma = 0.1
k0 = 50.
x0 = maxx/3.
U0 = 1./sqrt(sigma*sqrt(pi))*exp(-(x-x0)**2/(2*sigma**2)) * exp(1j*k0*x)
def Ut(t, x, x0=x0, k0=k0):
    p0 = hbar*k0
    norm = sqrt(sigma/(sqrt(pi)*(sigma**2+1j*hbar*t/m)))
    envelope = exp(-0.5*(x-x0-p0*t/m)**2/(sigma**2+1j*hbar*t/m))
    #phase = exp(1j/hbar*(p0*x-p0**2*t)/(2*m))
    phase = exp(1j*k0*(x-hbar*k0*t/(2*m)))
    return norm*envelope*phase


    

In [23]:

    

#plot(U0.real)
#plot(U0.imag)
plot(Ut(0, x).imag)
plot(Ut(0, x).real)
plot(abs(Ut(0, x)))


    

    
Out[23]:





[<matplotlib.lines.Line2D at 0x7f8b9d44f5c0>]






    

In [24]:

    

Ur[0, :] = Ut(0, x).real
Ui[0, :] = Ut(-0.5*dt, x).imag
#Ui[0, :] = U0.imag
schr_solve(Ur, Ui, dx, dt, nx, len(t))


    

In [9]:

    

#pcolormesh(sqrt(Ui**2+Ur**2)[:100,:], rasterized=True)
pcolormesh(Ui[:200,:], rasterized=True)


    

    
Out[9]:





<matplotlib.collections.QuadMesh at 0x7f8b9d070fd0>






    

In [10]:

    

# http://nbviewer.ipython.org/github/empet/Math/blob/master/DomainColoring.ipynb

def Hcomplex(z):# computes the hue corresponding to the complex number z
    H=np.angle(z)/(2*np.pi)+1
    return np.mod(H,1)
def g(x):
    return (1.0-1/(1+x**2))**0.2

U = Ur[:200,:] + 1j*Ui[:200,:]

from matplotlib.colors import hsv_to_rgb
H = Hcomplex(U)
V=g(abs(U)*0.5)
S = ones_like(H, float)

HSV = dstack((H,S,V))# V**0.2>V for V in[0,1];this choice  avoids too dark colors
RGB=hsv_to_rgb(HSV)
imshow(RGB)


    

    
Out[10]:





<matplotlib.image.AxesImage at 0x7f8b9cccef98>






    

In [11]:

    

from matplotlib import animation
from JSAnimation.IPython_display import display_animation


    

In [12]:

    

fig = figure()
ax = axes(xlim=(0,maxx),ylim=(-3,3))
line1, = ax.plot([],[],lw=1)
line2, = ax.plot([],[],lw=1)
line3, = ax.plot([],[],lw=2)
line4, = ax.plot([],[],lw=2)

def animate(data):
    line1.set_data(x,data[:len(x)])
    line2.set_data(x,data[len(x):])
    line3.set_data(x,sqrt(data[:len(x)]**2 + data[len(x):]**2))
    return line1,


    

In [13]:

    

Ui_int = 0.5*(Ui[:-1,:]+Ui[1:,:])
anim = animation.FuncAnimation(fig, animate, frames=hstack((Ui_int, Ur[:-1,:]))[:200:2,:], interval=100)
display_animation(anim, default_mode='loop')


    

    
Out[13]:








    
    

    
    

    –
    
    
    
    
    
    
    
    +
  

     Once 
     Loop 
     Reflect 
  

In [15]:

    

maxx = 5.
maxt = 0.08
nx = 500

sigma = 0.15
k0 = 40.
x0 = sigma*4.
E = (hbar**2/2.0/m)*(k0**2+0.5/sigma**2)

Ur, Ui, dx, dt, x, t = schr_init(maxx, maxt, E*1.5, nx, 0.7)

V = zeros_like(x)
xb = maxx*0.5
sb = 0.02
V = exp(-(x-xb)**2/(2*sb))*E*2.2
#V[(x>maxx*0.5) & (x<maxx*0.6)] = E*1.5

print(E)
print(len(t))
print(len(x))
Ur[0, :] = Ut(0, x).real
Ui[0, :] = Ut(-0.5*dt, x).imag
schr_solve(Ur, Ui, dx, dt, nx, len(t))


    

    




811.1111111111111
2434
502

In [16]:

    

Ui_int = 0.5*(Ui[:-1,:]+Ui[1:,:])
#ax.cla()
ax.set_xlim([0,maxx])
line4.set_data(x, (V-E)/V.max())
anim = animation.FuncAnimation(fig, animate, frames=hstack((Ui_int, Ur[:-1,:]))[::20,:], interval=100)
display_animation(anim, default_mode='loop')


    

    
Out[16]:








    
    

    
    

    –
    
    
    
    
    
    
    
    +
  

     Once 
     Loop 
     Reflect 
  

In [135]:

    

anim.save("animation.gif", writer='imagemagick', fps=10)


    

    




---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-135-f4608baa4b76> in <module>()
----> 1 anim.save("animation.gif", writer='imagemagick', fps=10)

/usr/lib/pymodules/python2.7/matplotlib/animation.pyc in save(self, filename, writer, fps, dpi, codec, bitrate, extra_args, metadata, extra_anim)
    597                 import warnings
    598                 warnings.warn("MovieWriter %s unavailable" % writer)
--> 599                 writer = writers.list()[0]
    600 
    601         verbose.report('Animation.save using %s' % type(writer),

IndexError: list index out of range





    




/usr/lib/pymodules/python2.7/matplotlib/animation.py:598: UserWarning: MovieWriter imagemagick unavailable
  warnings.warn("MovieWriter %s unavailable" % writer)

In [ ]: